import torch
from torch import nn
from einops import rearrange
from torch import einsum

class BidirectionalCrossAttentionSimple(nn.Module):
    """简化版双向交叉注意力，默认参数下无层归一化和talking heads"""
    def __init__(
        self,
        dim,                # 主序列维度
        heads = 8,          # 注意力头数
        dim_head = 64,      # 每个头的维度
        context_dim = None, # 上下文序列维度（默认与主序列相同）
        dropout = 0.,       # 统一dropout率
    ):
        super().__init__()
        # context_dim = default(context_dim, dim)
        
        self.heads = heads
        self.scale = dim_head ** -0.5
        inner_dim = dim_head * heads

        self.dropout = nn.Dropout(dropout)
        self.context_dropout = nn.Dropout(dropout)

        self.to_qk = nn.Linear(dim, inner_dim, bias = False)
        self.context_to_qk = nn.Linear(context_dim, inner_dim, bias = False)

        self.to_v = nn.Linear(dim, inner_dim, bias = False)
        self.context_to_v = nn.Linear(context_dim, inner_dim, bias = False)

        self.to_out = nn.Linear(inner_dim, dim)
        self.context_to_out = nn.Linear(inner_dim, context_dim)

    def forward(self, x, context, return_attn = False):
        
        # b - batch size（批大小）
        # i - 主序列长度（x的序列维度）
        # j - 上下文序列长度（context的序列维度）
        # h - attention heads数量（注意力头数）
        # b, i, j, h, device = x.shape[0], x.shape[-2], context.shape[-2], self.heads, x.device
        b, h, device = x.shape[0], self.heads, x.device

        # x = self.norm(x)
        # context = self.context_norm(context)

        # get shared query/keys and values for sequence and context

        qk, v = self.to_qk(x), self.to_v(x)
        context_qk, context_v = self.context_to_qk(context), self.context_to_v(context)

        # split out head
        qk, context_qk, v, context_v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (qk, context_qk, v, context_v))

        # get similarities

        sim = einsum('b h i d, b h j d -> b h i j', qk, context_qk) * self.scale

        # relative positional bias, if supplied

        # get attention along both sequence length and context length dimensions
        # shared similarity matrix

        attn = sim.softmax(dim = -1) # 主序列（查询序列）对上下文序列（键序列）的注意力分布
        context_attn = sim.softmax(dim = -2) # 上下文序列（查询序列）对主序列（键序列）的注意力分布

        # dropouts

        attn = self.dropout(attn)
        context_attn = self.context_dropout(context_attn)

        # src sequence aggregates values from context, context aggregates values from src sequence

        out = einsum('b h i j, b h j d -> b h i d', attn, context_v) 
        context_out = einsum('b h j i, b h j d -> b h i d', context_attn, v) 

        # merge heads and combine out
        out, context_out = map(lambda t: rearrange(t, 'b h n d -> b n (h d)'), (out, context_out))

        out = self.to_out(out)
        context_out = self.context_to_out(context_out)

        if return_attn:
            return out, context_out, attn, context_attn

        return out, context_out
    
if __name__ == "__main__":
    b_s = 16
    x_src = 1
    y_src = 1
    head = 1
    dim_head = 256
    n_feat = 128
    context_d = 128+29
    x = torch.randn(b_s, x_src, n_feat)
    context = torch.randn(b_s, y_src, context_d)
    model = BidirectionalCrossAttentionSimple(dim=n_feat, context_dim=context_d)
    output = model(x, context)